#include "StdInc.h"

#include "StdInc.h"
#include <atomic>

#include <Windows.h>
#include <shlobj.h>                 // For SHChangeNotify

#include <shell_default.h>
#include <shell_minlink.h>
#include <Reg.h>

// {C6162303-2F62-4E4B-81C9-B46D9D4BA37A}
const CLSID CLSID_IMGShellExtFolder = 
{ 0xc6162303, 0x2f62, 0x4e4b, { 0x81, 0xc9, 0xb4, 0x6d, 0x9d, 0x4b, 0xa3, 0x7a } };


void InitializeShellExtModule( void )
{
    // TODO.
}

void ShutdownShellExtModule( void )
{
    // TODO.
}

bool IsValidCLSID( REFCLSID rclsid )
{
    if ( rclsid == __uuidof(IShellFolder) )
    {
        return true;
    }

    return false;
}

struct shellfolder_objconstructor
{
    static inline HRESULT construct( IUnknown *pUnkOuter, REFIID riid, void *ppv )
    {
        if ( riid == __uuidof(IShellExtInit) )
        {


            return S_OK;
        }

        return E_NOINTERFACE;
    }
};

IClassFactory* CreateClassFactory( REFCLSID rclsid )
{
    return new (std::nothrow) shellClassFactory <shellfolder_objconstructor> ();
}

//
//   FUNCTION: DllRegisterServer
//
//   PURPOSE: Register the COM server and the thumbnail handler.
// 
STDAPI DllRegisterServer(void)
{
    HRESULT hr;

    wchar_t szModule[MAX_PATH];

    if (GetModuleFileNameW(module_instance, szModule, ARRAYSIZE(szModule)) == 0)
    {
        hr = HRESULT_FROM_WIN32(GetLastError());
        return hr;
    }

    RemoveUserChoiceFromProgram( L".img" );

    // Register the IMG shell view.
    hr = RegisterInprocServer(szModule, CLSID_IMGShellExtFolder, 
        L"imgext.IMGShellExtShellView Class", 
        L"Apartment");

    if ( SUCCEEDED(hr) )
    {
        hr = RegisterShellExtProgramOpenLink( L".img", L"IMGShellExtShellView" );

        if ( SUCCEEDED(hr) )
        {
            hr = RegisterShellExtShellFolder(L".img", CLSID_IMGShellExtFolder,
                SFGAO_BROWSABLE | SFGAO_FILESYSTEM | SFGAO_FOLDER | SFGAO_CANCOPY | SFGAO_CANRENAME | SFGAO_CANMOVE | SFGAO_STREAM);

            if ( SUCCEEDED( hr ) )
            {
                if ( auto SHChangeNotify = shell32min.SHChangeNotify )
                {
                    // This tells the shell to invalidate its file lists.
                    // There has been a change after all.
                    SHChangeNotify(SHCNE_ASSOCCHANGED, SHCNF_IDLIST, NULL, NULL);
                }
            }
        }
    }

    return hr;
}


//
//   FUNCTION: DllUnregisterServer
//
//   PURPOSE: Unregister the COM server and the thumbnail handler.
// 
STDAPI DllUnregisterServer(void)
{
    HRESULT hr = S_OK;

    wchar_t szModule[MAX_PATH];

    if (GetModuleFileNameW( module_instance, szModule, ARRAYSIZE(szModule) ) == 0)
    {
        hr = HRESULT_FROM_WIN32(GetLastError());
        return hr;
    }

    // Unregister the IMG shell view.
    if ( SUCCEEDED( hr ) )
    {
        hr = UnregisterInprocServer( CLSID_IMGShellExtFolder );

        if ( SUCCEEDED(hr) )
        {
            // Unregister the IMG handler.
            hr = UnregisterShellExtCommonDefault(L".img");

            if (SUCCEEDED(hr))
            {
                // Now delete the open link.
                hr = UnregisterShellExtProgramOpenLink(L".img", L"IMGShellExtShellView");
                
                if ( SUCCEEDED(hr) )
                {
                    // Should update our shell view.
                    if ( auto SHChangeNotify = shell32min.SHChangeNotify )
                    {
                        // This tells the shell to invalidate its file lists.
                        // There has been a change after all.
                        SHChangeNotify(SHCNE_ASSOCCHANGED, SHCNF_IDLIST, NULL, NULL);
                    }
                }
            }
        }
    }

    return hr;
}